#!/usr/bin/env python3
"""
1_extract_data.py
数据提取程序 - 从ak_a_hist_plus_fund.py提取原始数据
"""

import subprocess
import sys
import pandas as pd
import os
from datetime import datetime
import argparse

def extract_data(symbol, asset_type, start_date):
    """从ak_a_hist_plus_fund.py提取数据"""
    
    # 构建命令
    if asset_type == 'fund':
        cmd = [
            'python3', '../ak_a_hist_plus_fund.py',
            '--asset', 'openfund',
            '--symbol', symbol,
            '--start', start_date,
            '--end', 'today',
            '--export', '../csv_output/temp_raw.csv'
        ]
    elif asset_type == 'stock':
        cmd = [
            'python3', '../ak_a_hist_plus_fund.py',
            '--asset', 'stock',
            '--symbol', symbol,
            '--start', start_date,
            '--end', 'today',
            '--adjust', 'none',
            '--export', '../csv_output/temp_raw.csv'
        ]
    elif asset_type == 'etf':
        cmd = [
            'python3', '../ak_a_hist_plus_fund.py',
            '--asset', 'etf',
            '--symbol', symbol,
            '--start', start_date,
            '--end', 'today',
            '--adjust', 'none',
            '--export', '../csv_output/temp_raw.csv'
        ]
    
    # 执行命令
    try:
        result = subprocess.run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            print(f"Error: {result.stderr}", file=sys.stderr)
            return False
    except Exception as e:
        print(f"Error executing command: {e}", file=sys.stderr)
        return False
    
    # 读取临时文件并重命名列
    try:
        df = pd.read_csv('../csv_output/temp_raw.csv')
        
        # 标准化列名
        column_mapping = {
            '日期': 'date',
            '净值日期': 'date',
            '收盘': 'close',
            '开盘': 'open',
            '最高': 'high',
            '最低': 'low',
            '成交量': 'volume',
            '成交额': 'amount'
        }
        
        df = df.rename(columns=column_mapping)
        
        # 确保必要的列存在
        if 'close' not in df.columns:
            if '单位净值' in df.columns:
                df = df.rename(columns={'单位净值': 'close'})
                df['open'] = df['close']
                df['high'] = df['close']
                df['low'] = df['close']
                df['volume'] = 0
                df['amount'] = 0
        
        # 标准化日期格式
        df['date'] = pd.to_datetime(df['date'])
        df = df.sort_values('date')
        
        # 选择并排序列
        columns = ['date', 'open', 'high', 'low', 'close', 'volume', 'amount']
        for col in columns:
            if col not in df.columns:
                df[col] = 0 if col in ['volume', 'amount'] else df['close']
        
        df = df[columns]
        
        # 保存最终文件
        df.to_csv('../csv_output/level1_raw_data.csv', index=False)
        
        # 清理临时文件
        if os.path.exists('../csv_output/temp_raw.csv'):
            os.remove('../csv_output/temp_raw.csv')
            
        print(f"✅ 数据提取完成：{len(df)}条记录")
        return True
        
    except Exception as e:
        print(f"Error processing data: {e}", file=sys.stderr)
        return False

def main():
    parser = argparse.ArgumentParser(description='数据提取程序')
    parser.add_argument('--symbol', required=True, help='股票/基金代码')
    parser.add_argument('--type', choices=['stock', 'fund', 'etf'], required=True, help='资产类型')
    parser.add_argument('--start', required=True, help='开始日期(YYYY-MM-DD)')
    
    args = parser.parse_args()
    
    print(f"🔄 开始提取数据：{args.symbol} ({args.type})")
    
    success = extract_data(args.symbol, args.type, args.start)
    
    if success:
        print("✅ level1_raw_data.csv 已生成")
        sys.exit(0)
    else:
        print("❌ 数据提取失败")
        sys.exit(1)

if __name__ == "__main__":
    main()